#ifndef _LANCZOS_H_
#define _LANCZOS_H_
namespace SM
{
  //Lanczos algorithm
  template<typename MatrixAType>
    class LanczosIter
    {
    public:
      typedef typename MatrixAType::Scalar Scalar;
      typedef Matrix<Scalar,Dynamic, Dynamic> MatrixType;
      typedef Matrix<Scalar,Dynamic, 1> VectorType;
      
    LanczosIter(const MatrixAType &_A,const VectorType &_V0):A(_A),V0(_V0),dim(0),totaldim(_A.rows()){setV0(_V0);}
    LanczosIter(const MatrixAType &_A):A(_A),dim(0),totaldim(_A.rows()){
	srand((unsigned int) time(0)); 
	V0=VectorType::Random(_A.rows());
	//	V0=MatrixType::Ones(_A.rows(),1);//another choice
	V0/=sqrt(V0.transpose()*V0);
      }
      void clear();
      void setV0(const VectorType & _V0)
      {
	V0=_V0;
	V0/=sqrt(V0.transpose()*V0);
	clear();
      }
      bool next();//get next lanczos vector
      int getAlphaBetaQ(VectorType & diag, VectorType & subdiag,MatrixType & Q,int _dim,bool yesvecs=true);//return the tridiagnal matrix in form of diag and subdiag, every coloum of Q is a Lanczos vector
      vector<Scalar> alpha;
      vector<Scalar> beta;
      vector<VectorType> V;
      VectorType u;//the rest
      int dim;


      const int totaldim;
      const MatrixAType & A;
      VectorType V0;
    };

  template<typename MatrixAType>
    void LanczosIter<MatrixAType>::clear()
    {
      alpha.clear();
      beta.clear();
      V.clear();
      dim=0;
    }

  template<typename MatrixAType>
    bool LanczosIter<MatrixAType>::next()
    {
      if(alpha.empty())
	{
	  V.push_back(V0);
	  u=A*V[0];
	  alpha.push_back( u.transpose()*V[0] );
	  u-=alpha[0]*V[0];
	}
      else
	{
	  Scalar norm = sqrt(u.transpose()*u);
	  if(abs(norm)<1e-8)
	    {
	      return false;
	    }
	  else
	    {
	      beta.push_back( norm );
	      V.push_back( u/norm );

	      int Index=alpha.size();
	      u=A*V[Index];
	      alpha.push_back( u.transpose()*V[Index] );
	      // u-=alpha[ Index ]*V[ Index ] + beta[ Index -1]*V[ Index -1];

	      //reorthogonal in case of numeric error accumulating
	      for(int i=Index;i>=0;i--)
		{
		  Scalar t=u.transpose()*V[i];
		  u-=t*V[i];
		}
	      
	    }
	}
      dim=alpha.size();
      return true;
    }

  template<typename MatrixAType>
    int LanczosIter<MatrixAType>::getAlphaBetaQ(VectorType & diag, VectorType & subdiag,MatrixType & Q,int _dim,bool yesvecs)
    {
      if(dim<_dim)
	{
	  while(next())
	    {
	      if(dim==_dim)
		break;
	    }
	}
      if(dim<_dim)
	{
	  _dim=dim;
	  std::cout<<"warning: The lanczos iteration terminated, you may need to change the starting vec.\n";
	}
      diag.resize(_dim);
      subdiag.resize(_dim);
      for(int i=0;i<_dim;i++)
	{
	  diag[i]=alpha[i];
	}
      for(int i=0;i<_dim-1;i++)
	{
	  subdiag[i]=beta[i];
	}
      if(yesvecs)
	{
	  Q.resize( A.rows(),_dim );
	  for(int i=0;i<_dim;i++)
	    {
	      Q.col(i)=V[i];
	    }
	}
      return _dim;
    }



  //****************************************************************
  //sort algorithm, sort according to eigvals and reorder the eigvecs respectively
  template<typename Scalar>
    void quickSort(Matrix<Scalar,Dynamic,1> & eigvals,Matrix<Scalar,Dynamic,Dynamic> &eigvecs,int startPos, int endPos) 
    { 
      int i,j; 
      Scalar key; 
      key=eigvals[startPos];
      Matrix<Scalar,Dynamic,1> Vec=eigvecs.col(startPos);
      i=startPos; 
      j=endPos;
      while(i<j) 
	{ 
	  while(!(eigvals[j]<key)&& i<j) --j; 
	  eigvals[i]=eigvals[j];
	  eigvecs.col(i)=eigvecs.col(j);
	  while(eigvals[i]<key && i<j) ++i; 
	  eigvals[j]=eigvals[i];
	  eigvecs.col(j)=eigvecs.col(i);
	} 
      eigvals[i]=key;
      eigvecs.col(i)=Vec; 
      if(i-1>startPos) quickSort(eigvals,eigvecs,startPos,i-1); 
      if(endPos>i+1) quickSort(eigvals,eigvecs,i+1,endPos); 
    } 

  template<typename Scalar>
    void quickSort(Matrix<Scalar,Dynamic,1> & eigvals,int startPos, int endPos) 
    { 
      int i,j; 
      Scalar key; 
      key=eigvals[startPos];
      i=startPos; 
      j=endPos;
      while(i<j) 
	{ 
	  while(!(eigvals[j]<key)&& i<j) --j; 
	  eigvals[i]=eigvals[j];
	  while(eigvals[i]<key && i<j) ++i; 
	  eigvals[j]=eigvals[i];
	} 
      eigvals[i]=key;
      if(i-1>startPos) quickSort(eigvals,startPos,i-1); 
      if(endPos>i+1) quickSort(eigvals,i+1,endPos); 
    } 

  
  //****************************************************************
  template<typename MatrixAType>
    class LanczosEigenSolver
    {
    public:
      typedef typename MatrixAType::Scalar Scalar;
      typedef Matrix<Scalar, Dynamic, Dynamic> MatrixType;
      typedef Matrix<Scalar, Dynamic, 1> VectorType;

    LanczosEigenSolver(LanczosIter<MatrixAType> & _lanczositer):lanczositer(_lanczositer){}

      void setV0(const VectorType & V0)
      {
	lanczositer.setV0(V0);
      }
      void compute(int _max_dim=-1,bool yesvecs=true);
      const MatrixType& eigenvectors() const
      {
	return m_eivec;
      }
      const VectorType& eigenvalues() const
      {
	return m_eivalues;
      }
    private:
      LanczosIter<MatrixAType> & lanczositer;
      MatrixType m_eivec;
      VectorType m_eivalues;
    };

  template<typename MatrixAType>
    void LanczosEigenSolver<MatrixAType>::compute(int max_dim,bool yesvecs)
    {
      if(max_dim==0) return;
      if(max_dim>lanczositer.totaldim||max_dim==-1) max_dim=lanczositer.totaldim;
 
      VectorType &diag=m_eivalues;
      VectorType subdiag;
      int n=lanczositer.getAlphaBetaQ(diag,subdiag,m_eivec,max_dim,yesvecs);
      double scale=abs(diag[0]);
      for(int i=0;i<n;i++)
	{
	  double temp=abs(diag[i]);
	  if(scale<temp)
	    scale=temp;
	  if(i<n-1)	
	    {
	      temp=abs(subdiag[i]);
	      if(scale<temp)
		scale=temp;
	    }
	}

      if(scale==double(0)) scale = double(1);
      diag/=scale;
      subdiag/=scale;

      TriDiagQL(diag,subdiag,m_eivec,n,yesvecs);
  
  
      // Sort eigenvalues and corresponding vectors.
      // TODO make the sort optional ?
      // TODO use a better sort algorithm !!

      if(yesvecs)
	quickSort(m_eivalues,m_eivec,0,n-1);
      else
	quickSort(m_eivalues,0,n-1);
  
      // scale back the eigen values
      m_eivalues *= scale;
    }
}


#endif
